import torch
import torch.nn as nn
import torch.nn.functional as F


class SimplifiedFusion(nn.Module):
    """简化版特征融合模块"""

    def __init__(self, in_ch=256):
        super().__init__()
        self.channel_adjust = nn.Conv2d(in_ch * 2, in_ch, 1)
        self.attention = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_ch, in_ch // 8, 1),
            nn.ReLU(),
            nn.Conv2d(in_ch // 8, in_ch, 1),
            nn.Sigmoid()
        )

    def forward(self, f1, f2):
        # 确保特征图尺寸一致
        if f1.shape[-2:] != f2.shape[-2:]:
            f2 = F.interpolate(f2, size=f1.shape[-2:], mode='bilinear', align_corners=True)

        fused = torch.cat([f1, f2], dim=1)
        fused = self.channel_adjust(fused)  # [b,256,h,w]
        return fused * self.attention(fused)


class RobustMultiModalNet(nn.Module):
    def __init__(self,channel=1,lei=4):
        super().__init__()

        # 统一特征提取器
        def build_stream():
            return nn.Sequential(
                nn.Conv2d(3, 64, 3, padding=1),
                nn.MaxPool2d(2),  # 128x128
                nn.BatchNorm2d(64),
                nn.ReLU(),

                nn.Conv2d(64, 128, 3, padding=1),
                nn.MaxPool2d(2),  # 64x64
                nn.BatchNorm2d(128),
                nn.ReLU(),

                nn.Conv2d(128, 256, 3, padding=1),
                nn.MaxPool2d(2),  # 32x32
                nn.BatchNorm2d(256),
                nn.ReLU(),
            )

        self.stream1 = build_stream()
        self.stream2 = build_stream()

        # 非图像处理
        self.mlp = nn.Sequential(
            nn.Linear(42, 128),
            nn.ReLU(),
            nn.Linear(128, 256)
        )

        # 融合模块
        self.fusion = SimplifiedFusion()

        # 分类头
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(256, 4)
        )

    def forward(self, x1, x2, x3):
        # 图像特征提取 (统一输出256通道)
        f1 = self.stream1(x1)  # [b,256,32,32]
        f2 = self.stream2(x2)  # [b,256,32,32]

        # 向量特征处理
        vec_feat = self.mlp(x3).unsqueeze(-1).unsqueeze(-1)  # [b,256,1,1]

        # 特征融合
        fused = self.fusion(f1, f2)

        # 加入向量特征
        output = fused * vec_feat
        return self.classifier(output)

    # 验证测试


if __name__ == "__main__":
    model = RobustMultiModalNet()
    x1 = torch.randn(2, 3, 256, 256)
    x2 = torch.randn(2, 3, 256, 256)
    x3 = torch.randn(2, 42)

    try:
        with torch.no_grad():
            output = model(x1, x2, x3)
            print(f"验证通过！输出形状: {output.shape}")  # 预期输出: torch.Size([2, 4])
    except Exception as e:
        print("错误信息:", e)